
# 路径
TRAINING_DATA_PATH = "./datasets/chemblv31/train_30.txt"  # 用于微调的30%数据
FILTERED_TRAINING_DATA_PATH = "./datasets/chemblv31/train_30_filtered.txt"  # 筛选后的训练集保存路径

MAX_LENGTH = 50  # 限制生成序列的最大长度


def filter_smiles_by_length(input_file, output_file, max_length):
    """
    筛选SMILES字符串，保留长度不超过max_length的分子。
    :param input_file: 输入的SMILES数据文件路径
    :param output_file: 输出的筛选后的SMILES数据文件路径
    :param max_length: SMILES字符串的最大长度
    """
    valid_smiles = []
    with open(input_file, "r") as f:
        for line in f:
            smile = line.strip()
            if len(smile) <= max_length:
                valid_smiles.append(smile)

    # 保存筛选后的SMILES数据
    with open(output_file, "w") as f:
        for smile in valid_smiles:
            f.write(smile + "\n")

    print(f"Filtered data saved to {output_file}, {len(valid_smiles)} valid molecules found.")


# 筛选训练数据
filter_smiles_by_length(TRAINING_DATA_PATH, FILTERED_TRAINING_DATA_PATH, MAX_LENGTH)

